cmake_minimum_required(VERSION 3.22)

option(WITH_LIBTORCH "build _YOLO_LIBTORCH" ON)
option(WITH_ONNXRUNTIME "build _YOLO_ONNXRUNTIME" ON)
option(WITH_OPENCV "build _YOLO_OPENCV" ON)
option(WITH_OPENVINO "build _YOLO_OPENVINO" ON)
option(WITH_TENSORRT "build _YOLO_TENSORRT" ON)
option(WITH_CUDA "build _WITH_CUDA" ON)
option(WITH_CUDA_PREPROCESS "build _CUDA_PREPROCESS" ON)
option(WITH_CUDA_POSTPROCESS "build _CUDA_POSTPROCESS" ON)

if(WITH_CUDA)
    set(CMAKE_CUDA_COMPILER "/usr/local/cuda-12.8/bin/nvcc")
    project(yolo LANGUAGES C CXX CUDA)
else()
    project(yolo LANGUAGES C CXX)
endif()

file(GLOB SRCS ${CMAKE_CURRENT_SOURCE_DIR}/*.h ${CMAKE_CURRENT_SOURCE_DIR}/*.cpp)

include_directories(${CMAKE_CURRENT_SOURCE_DIR})

add_compile_options(-w)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_BUILD_TYPE "Release")

if(WITH_LIBTORCH)
    add_compile_options(-D_YOLO_LIBTORCH)
    aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/libtorch SRCS)
    include_directories(${CMAKE_CURRENT_SOURCE_DIR}/libtorch)
endif()

if(WITH_ONNXRUNTIME)
    add_compile_options(-D_YOLO_ONNXRUNTIME)
    aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime SRCS)
    include_directories(${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime)
endif()

if(WITH_OPENCV)
    add_compile_options(-D_YOLO_OPENCV)
    aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/opencv SRCS)
    include_directories(${CMAKE_CURRENT_SOURCE_DIR}/opencv)
endif()

if(WITH_OPENVINO)
    add_compile_options(-D_YOLO_OPENVINO)
    aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/openvino SRCS)
    include_directories(${CMAKE_CURRENT_SOURCE_DIR}/openvino)
endif()

if(WITH_CUDA AND WITH_TENSORRT)
    add_compile_options(-D_YOLO_TENSORRT)
    list(APPEND SRCS ${CMAKE_CURRENT_SOURCE_DIR}/tensorrt/yolo_tensorrt.h ${CMAKE_CURRENT_SOURCE_DIR}/tensorrt/yolo_tensorrt.cpp)
    include_directories(${CMAKE_CURRENT_SOURCE_DIR}/tensorrt)
    if(WITH_CUDA_PREPROCESS)
        list(APPEND SRCS ${CMAKE_CURRENT_SOURCE_DIR}/tensorrt/preprocess.cuh ${CMAKE_CURRENT_SOURCE_DIR}/tensorrt/preprocess.cu)
        add_compile_options(-D_CUDA_PREPROCESS)
        if(WITH_CUDA_POSTPROCESS)
            list(APPEND SRCS ${CMAKE_CURRENT_SOURCE_DIR}/tensorrt/decode.cuh ${CMAKE_CURRENT_SOURCE_DIR}/tensorrt/decode.cu)
            add_compile_options(-D_CUDA_POSTPROCESS)
        endif()
    endif()
endif()

if(CMAKE_SYSTEM_NAME MATCHES "Windows")   
    add_executable(yolo ${SRCS})

    if(WITH_CUDA)
        find_package(CUDAToolkit REQUIRED)
        include_directories(${CUDA_INCLUDE_DIRS})
        link_directories(${CUDA_LIB_PATH})
        target_link_libraries(yolo CUDA::cudart)
    endif()

    set(OPENCV_INCLUDE_DIRS D:/opencv-4.10.0/install/include)
    set(OPENCV_LIBRARY_DIRS D:/opencv-4.10.0/install/x64/vc17/lib)
    include_directories(${OPENCV_INCLUDE_DIRS})
    link_directories(${OPENCV_LIBRARY_DIRS})
    target_link_libraries(yolo ${OPENCV_LIBRARY_DIRS}/*.lib)

    if(WITH_LIBTORCH)
        set(LIBTORCH_INCLUDE_DIRS D:/libtorch-win-shared-with-deps-2.5.0+cu124/include D:/libtorch-win-shared-with-deps-2.5.0+cu124/include/torch/csrc/api/include)
        set(LIBTORCH_LIBRARY_DIRS D:/libtorch-win-shared-with-deps-2.5.0+cu124/lib)
        include_directories(${LIBTORCH_INCLUDE_DIRS})
        link_directories(${LIBTORCH_LIBRARY_DIRS})
        target_link_libraries(yolo ${LIBTORCH_LIBRARY_DIRS}/*.lib)
        target_link_options(yolo PRIVATE "/INCLUDE:\"?ignore_this_library_placeholder@@YAHXZ\"")
    endif()

    if(WITH_ONNXRUNTIME)
        set(ONNXRUNTIME_INCLUDE_DIRS D:/onnxruntime-win-x64-gpu-1.20.0/include)
        set(ONNXRUNTIME_LIBRARY_DIRS D:/onnxruntime-win-x64-gpu-1.20.0/lib)
        include_directories(${ONNXRUNTIME_INCLUDE_DIRS})
        link_directories(${ONNXRUNTIME_LIBRARY_DIRS})
        target_link_libraries(yolo ${ONNXRUNTIME_LIBRARY_DIRS}/*.lib)
    endif()

    if(WITH_OPENVINO)
        set(OPENVINO_INCLUDE_DIRS D:/w_openvino_toolkit_windows_2024.4.0.16579.c3152d32c9c_x86_64/runtime/include)
        set(OPENVINO_LIBRARY_DIRS D:/w_openvino_toolkit_windows_2024.4.0.16579.c3152d32c9c_x86_64/runtime/lib/intel64/Release)
        include_directories(${OPENVINO_INCLUDE_DIRS})
        link_directories(${OPENVINO_LIBRARY_DIRS})
        target_link_libraries(yolo ${OPENVINO_LIBRARY_DIRS}/*.lib)
    endif()

    if(WITH_TENSORRT)
	set(TENSORRT_INCLUDE_DIRS D:/TensorRT-10.6.0.26/include)
	set(TENSORRT_LIBRARY_DIRS D:/TensorRT-10.6.0.26/lib)
        include_directories(${TENSORRT_INCLUDE_DIRS})
        link_directories(${TENSORRT_LIBRARY_DIRS})
        target_link_libraries(yolo ${TENSORRT_LIBRARY_DIRS}/*.lib)
    endif()
endif()

if(CMAKE_SYSTEM_NAME MATCHES "Linux")
    if(WITH_CUDA)
        find_package(CUDA REQUIRED)
        include_directories(${CUDA_INCLUDE_DIRS})
        cuda_add_executable(yolo ${SRCS})
        target_link_libraries(yolo ${CUDA_LIBRARIES})
    else()
        add_executable(yolo ${SRCS})
    endif()

    find_package(OpenCV REQUIRED)
    include_directories(${OpenCV_INCLUDE_DIRS} ${OpenCV2_INCLUDE_DIRS})
    target_link_libraries(yolo ${OpenCV_LIBRARIES})  

    if(WITH_LIBTORCH)
        set(Torch_DIR /workspace/libtorch/share/cmake/Torch)
        find_package(Torch REQUIRED)
        include_directories(${TORCH_INCLUDE_DIRS})
        target_link_libraries(yolo ${TORCH_LIBRARIES})
    endif()

    if(WITH_ONNXRUNTIME)
        set(ONNXRUNTIME_INCLUDE_DIRS /workspace/onnxruntime-linux-x64-gpu-1.22.0/include)
        set(ONNXRUNTIME_LIBRARY_DIRS /workspace/onnxruntime-linux-x64-gpu-1.22.0/lib)
        include_directories(${ONNXRUNTIME_INCLUDE_DIRS})
        target_link_libraries(yolo ${ONNXRUNTIME_LIBRARY_DIRS}/libonnxruntime.so)
    endif()

    if(WITH_OPENVINO)
        set(OPENVINO_INCLUDE_DIRS /workspace/openvino_toolkit_ubuntu24_2025.2.0.19140.c01cd93e24d_x86_64/runtime/include)
        set(OPENVINO_LIBRARY_DIRS /workspace/openvino_toolkit_ubuntu24_2025.2.0.19140.c01cd93e24d_x86_64/runtime/lib/intel64)
        include_directories(${OPENVINO_INCLUDE_DIRS})
        target_link_libraries(yolo ${OPENVINO_LIBRARY_DIRS}/libopenvino.so ${OPENVINO_LIBRARY_DIRS}/libopenvino_intel_cpu_plugin.so ${OPENVINO_LIBRARY_DIRS}/libopenvino_intel_gpu_plugin.so)
    endif()

    if(WITH_TENSORRT)
        set(TENSORRT_INCLUDE_DIRS /usr/include/x86_64-linux-gnu)
        set(TENSORRT_LIBRARY_DIRS /usr/lib/x86_64-linux-gnu)
        include_directories(${TENSORRT_INCLUDE_DIRS})
        link_directories(${TENSORRT_LIBRARY_DIRS})
        target_link_libraries(yolo ${TENSORRT_LIBRARY_DIRS}/libnvinfer.so)
    endif() 
endif()
